import QuantLib as ql
import pandas as pd
import seaborn as sb
import math
from WindPy import w
import matplotlib.pyplot as plt
import datetime
import numpy as np
import Core.Config as Config
import Core.Gadget as Gadget

# Wind EDB Function
def EDB(name, code, datetime1, datetime2):
    #
    if datetime1 == None:
        datetime1 = datetime.datetime(2000,1,1)
    #
    mydata = w.edb(code, datetime1, datetime2, 'Fill=Previous')
    mydata_df = pd.DataFrame(data=mydata.Data[0], index=mydata.Times)
    mydata_df.columns = [name]
    return mydata_df


def Save_TimeSeries_ToDatabase(database, df, database_name, table_name):
    #
    index = df.index.values
    # 只针对第一列进行操作
    factorName = df.columns[0]
    df = df[factorName]
    #
    print("Saving", factorName)
    # print(df)

    newDocuments = []
    #
    for i in range(len(df)):
        name = factorName
        # print('name\n', name)
        date = index[i]
        #
        value = df.iloc[i]
        if np.isnan(value):
            continue
        #
        newDocument = {}
        newDocument["Name"] = name
        newDocument["Value"] = float(value)
        newDocument["Date"] = date
        newDocument["DateTime"] = date
        newDocument["Key2"] = newDocument["Name"] + "_" + date.strftime('%Y-%m-%d')
        newDocuments.append(newDocument)
    #
    database.Upsert_Many(database_name, table_name, [], newDocuments)
    print("Saved #", len(newDocuments), factorName)


def Build_TermStructure_Curve(reference_date, terms, rates, plot=False):
    dates = []
    dfs = []
    #
    cal = ql.China(ql.China.IB)
    today = ql.Date(reference_date.day, reference_date.month, reference_date.year)
    dc = ql.ActualActual(ql.ActualActual.Actual365)
    #
    settlementDays = 0
    # settlement = cal.advance(today, settlementDays, ql.Days)
    #
    for i in range(len(terms)):
        dates.append(today + ql.Period(terms[i], ql.Months))
        rate = rates[i]
        dfs.append(rate / 100.0)
    #
    curve = ql.ZeroCurve(dates, dfs, dc, cal, ql.Linear(), ql.Compounded, ql.Annual)

    #
    if plot == False:
        return curve

    data = []
    for i in range(200):
        tmpDate = today + ql.Period(i, ql.Days)
        rate = 100.0 * curve.zeroRate(tmpDate, dc, ql.Compounded, ql.Annual).rate()
        # print(tmpDate, rate)
        data.append([i, rate])
    #
    df = pd.DataFrame(data=data, columns=["Date","Rate"])
    df.plot(x="Date", y=["Rate"], grid=True)
    plt.show()
    return curve


def Download_TBond(database, datetime1, datetime2):
    #
    df = EDB("TBond_YTM_0", "M1004136", datetime1, datetime2)
    Save_TimeSeries_ToDatabase(database, df, "Bond", "TBond")
    #
    df = EDB("TBond_YTM_1M", "M1004677", datetime1, datetime2)
    Save_TimeSeries_ToDatabase(database, df, "Bond", "TBond")
    #
    df = EDB("TBond_YTM_2M", "M1004829", datetime1, datetime2)
    Save_TimeSeries_ToDatabase(database, df, "Bond", "TBond")
    #
    df = EDB("TBond_YTM_3M", "S0059741", datetime1, datetime2)
    Save_TimeSeries_ToDatabase(database, df, "Bond", "TBond")
    #
    df = EDB("TBond_YTM_6M", "S0059742", datetime1, datetime2)
    Save_TimeSeries_ToDatabase(database, df, "Bond", "TBond")
    #
    df = EDB("TBond_YTM_9M", "S0059743", datetime1, datetime2)
    Save_TimeSeries_ToDatabase(database, df, "Bond", "TBond")
    #
    df = EDB("TBond_YTM_1Y", "S0059744", datetime1, datetime2)
    Save_TimeSeries_ToDatabase(database, df, "Bond", "TBond")
    #
    df = EDB("TBond_YTM_2Y", "S0059745", datetime1, datetime2)
    Save_TimeSeries_ToDatabase(database, df, "Bond", "TBond")


def Download_TBond_Automatic(database, datetime2):
    #
    documents = database.Find('Bond', 'TBond', filter={"Name": "TBond_YTM_0"}, sort=[("Date", 1)])
    max_date = documents[-1]["Date"]
    #
    Download_TBond(database, max_date, datetime2.date())


# 获取关键期限数据
def Find_TBond_TermStructure_Rate(database, datetime1, datetime2):
    data = {}
    #
    names = ["TBond_YTM_0", "TBond_YTM_1M", "TBond_YTM_2M", "TBond_YTM_3M",
             "TBond_YTM_6M", "TBond_YTM_9M", "TBond_YTM_1Y", "TBond_YTM_2Y"]
    #
    if datetime1 == None:
        datetime1 = datetime2 + datetime.timedelta(days=-30)
    #
    filter = {}
    filter["Date"] = {">=": datetime1, "<=": datetime2}

    # ---Find in database---
    df_tbond = pd.DataFrame()
    for name in names:
        filter["Name"] = name
        documents = database.Find('Bond', 'TBond', filter=filter, sort=[("Date", 1)])
        df = Gadget.DocumentsToDataFrame(documents, keep=["Date", "Value"], rename={"Value": name})
        #
        if df_tbond.empty:
            df_tbond = df
        else:
            df_tbond = pd.merge(df_tbond, df, on="Date")

    # ---Export---
    recent = df_tbond.iloc[-1]
    rates = df_tbond.mean()
    # print(rates)
    data["Date"] = recent["Date"]
    data["Terms"] = [0, 1, 2, 3, 6, 9, 12, 24]
    data["Rates"] = [rates["TBond_YTM_0"], rates["TBond_YTM_1M"], rates["TBond_YTM_2M"], rates["TBond_YTM_3M"],
                     rates["TBond_YTM_6M"], rates["TBond_YTM_9M"], rates["TBond_YTM_1Y"], rates["TBond_YTM_2Y"]]
    #
    return data


# 建立国债收益率期限结构
def Build_TBond_TermStructure_Curve(database, datetime1, datetime2):
    data = Find_TBond_TermStructure_Rate(database, datetime1, datetime2)
    terms = data["Terms"]
    rates = data["Rates"]
    reference_date = data["Date"]
    curve = Build_TermStructure_Curve(reference_date, terms, rates, plot=False)
    return curve

def testingYields2():
    dates = []
    dfs = []

    cal = ql.China(ql.China.IB)
    today = ql.Date(23, ql.July, 2018)
    dc = ql.ActualActual(ql.ActualActual.ISMA)

    settlementDays = 0
    settlement = cal.advance(
        today, settlementDays, ql.Days)

    dates.append(settlement)
    dates.append(settlement + ql.Period(1, ql.Years))
    dates.append(settlement + ql.Period(2, ql.Years))
    dates.append(settlement + ql.Period(3, ql.Years))
    dates.append(settlement + ql.Period(4, ql.Years))
    dates.append(settlement + ql.Period(5, ql.Years))
    dates.append(settlement + ql.Period(6, ql.Years))
    dates.append(settlement + ql.Period(7, ql.Years))
    dates.append(settlement + ql.Period(8, ql.Years))
    dates.append(settlement + ql.Period(9, ql.Years))
    dates.append(settlement + ql.Period(10, ql.Years))
    dates.append(settlement + ql.Period(15, ql.Years))
    dates.append(settlement + ql.Period(20, ql.Years))
    dates.append(settlement + ql.Period(30, ql.Years))

    dfs.append(0.0000 / 100.0)
    dfs.append(3.0544 / 100.0)
    dfs.append(3.1565 / 100.0)
    dfs.append(3.2531 / 100.0)
    dfs.append(3.2744 / 100.0)
    dfs.append(3.2964 / 100.0)
    dfs.append(3.4092 / 100.0)
    dfs.append(3.5237 / 100.0)
    dfs.append(3.5264 / 100.0)
    dfs.append(3.5298 / 100.0)
    dfs.append(3.5337 / 100.0)
    dfs.append(3.8517 / 100.0)
    dfs.append(3.8884 / 100.0)
    dfs.append(4.0943 / 100.0)

    tmpDate1 = settlement + ql.Period(7, ql.Years)
    tmpDate2 = settlement + ql.Period(8, ql.Years)
    curve = ql.ZeroCurve(
        dates, dfs, dc, cal, ql.Linear(), ql.Compounded, ql.Annual)

    print(
        "零息收益率：",
        curve.zeroRate(
            tmpDate2, dc, ql.Compounded, ql.Annual))
    print(
        "贴现因子：",
        curve.discount(tmpDate2))
    print(
        "7Y - 8Y 远期收益率：",
        curve.forwardRate(
            tmpDate1, tmpDate2, dc, ql.Compounded, ql.Annual))


def testingYields3():
    maturities = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 20, 30, 50]
    pars = [0.030544, 0.031549, 0.032489, 0.032702, 0.032915, 0.033958, 0.03500,
            0.035050, 0.035100, 0.035150, 0.037765, 0.038163, 0.039568, 0.03972]
    spots = [0.030544, 0.031565, 0.032531, 0.032744, 0.032964, 0.034092, 0.035237,
             0.035264, 0.035298, 0.035337, 0.038517, 0.038884, 0.040943, 0.040720]

    numberOfBonds = len(maturities)
    cleanPrice = [100.0] * numberOfBonds
    quote = [ql.SimpleQuote(c) for c in cleanPrice]
    # quoteHandle = [ql.QuoteHandle()] * numberOfBonds
    # for i in range(len(quoteHandle)):
    #     quoteHandle[i].linkTo(quote[i])
    quoteHandle = [ql.QuoteHandle(q) for q in quote]

    frequency = ql.Annual
    dc = ql.ActualActual(ql.ActualActual.ISMA)
    accrualConvention = ql.ModifiedFollowing
    convention = ql.ModifiedFollowing
    redemption = 100.0
    calendar = ql.China(ql.China.IB)

    today = calendar.adjust(ql.Date(23, 7, 2018))
    # ql.Settings.evaluationDate = today
    ql.Settings.instance().evaluationDate = today

    bondSettlementDays = 0
    bondSettlementDate = calendar.advance(
        today,
        ql.Period(bondSettlementDays, ql.Days))

    instruments = []

    for j in range(len(maturities)):
        maturity = calendar.advance(
            bondSettlementDate,
            ql.Period(maturities[j], ql.Years))
        schedule = ql.Schedule(
            bondSettlementDate,
            maturity,
            ql.Period(frequency),
            calendar,
            accrualConvention,
            accrualConvention,
            ql.DateGeneration.Backward,
            False)

        helper = ql.FixedRateBondHelper(
            quoteHandle[j],
            bondSettlementDays,
            100.0,
            schedule,
            [pars[j]],
            dc,
            convention,
            redemption)

        instruments.append(helper)

    tolerance = 1.0e-10
    max = 5000
    svensson = ql.SvenssonFitting()

    ts0 = ql.PiecewiseLogCubicDiscount(
        bondSettlementDate,
        instruments,
        dc)
    ts1 = ql.FittedBondDiscountCurve(
        bondSettlementDate,
        instruments,
        dc,
        svensson,
        tolerance,
        max)

    spline = []
    sv = []

    print('{0:>9}{1:>9}{2:>9}{3:>9}'.format(
        "tenor", "spot", "spline", "svensson"))

    for i in range(len(instruments)):
        cfs = instruments[i].bond().cashflows()
        cfSize = len(instruments[i].bond().cashflows())

        tenor = dc.yearFraction(today, cfs[cfSize - 1].date())

        print(
            '{0:9.3f}{1:9.3f}{2:9.3f}{3:9.3f}'.format(
                tenor,
                100.0 * spots[i],
                100.0 * ts0.zeroRate(cfs[cfSize - 1].date(), dc, ql.Compounded, frequency).rate(),
                100.0 * ts1.zeroRate(cfs[cfSize - 1].date(), dc, ql.Compounded, frequency).rate()))

        spline.append(ts0.zeroRate(cfs[cfSize - 1].date(), dc, ql.Compounded, frequency).rate())
        sv.append(ts1.zeroRate(cfs[cfSize - 1].date(), dc, ql.Compounded, frequency).rate())

    df = pd.DataFrame(
        dict(
            mat=maturities * 4,
            rate=pars + spots + spline + sv,
            type=['par'] * 14 + ['spot'] * 14 + ['spline'] * 14 + ['sv'] * 14))

    print(ts1.fitResults().solution())
    return df


if __name__ == '__main__':
    #
    config = Config.Config()
    database = config.DataBase("MySQL")

    #
    datetime_now = datetime.datetime.now()
    w.start()
    # ---Cold Start---
    # datetime1 = datetime.datetime(2000, 1, 1)
    # datetime2 = datetime.datetime(2020, 1, 1)
    # Download_TBond(database, datetime1, datetime2)

    # 日常更新
    # Download_TBond_Automatic(database, datetime_now)

    datetime2 = datetime.datetime(2019, 12, 31)
    # Find_TBond_TermStructure_Rate(database, datetime1=None, datetime2)
    Build_TBond_TermStructure_Curve(database, datetime1=None, datetime2=datetime2)

    # reference_date = datetime.datetime(2019,1,1)
    # terms = [0, 3, 6, 9, 12, 24]
    # rates = [1.4197, 2.0373, 2.1576, 2.3061, 2.4264, 2.5006]
    # Build_TermStructure(reference_date, terms, rates, plot=True)

    # testingYields2()
    #
    # rs = testingYields3()
    # #
    # sb.relplot(
    #     x='mat', y='rate', kind='line', hue='type',
    #     data=rs, height=5, aspect=1.6)